"""
联邦学习API路由
"""

from fastapi import APIRouter, HTTPException
from pydantic import BaseModel
from typing import List, Dict, Any, Optional
from src.services.federated_learning_service import FederatedLearningIntegration
from src.research_core.model_manager import ModelType

router = APIRouter(prefix="/federated", tags=["federated-learning"])

# 初始化联邦学习集成实例
federated_integration = FederatedLearningIntegration()

class FederatedNodeConfig(BaseModel):
    node_id: str
    model_type: str  # ModelType的字符串表示
    data_loader: str  # 数据加载器描述

class FederatedTrainingRequest(BaseModel):
    model_type: str
    rounds: int = 10
    node_configs: List[FederatedNodeConfig]

@router.post("/setup")
async def setup_federated_training(request: FederatedTrainingRequest):
    """设置联邦训练"""
    try:
        # 转换model_type字符串为ModelType枚举
        model_type = ModelType(request.model_type)
        
        # 转换节点配置
        node_configs = []
        for config in request.node_configs:
            node_configs.append({
                "node_id": config.node_id,
                "model_type": ModelType(config.model_type),
                "data_loader": config.data_loader  # 实际实现中需要解析为实际函数
            })
        
        await federated_integration.setup_federated_training(node_configs)
        return {"message": "联邦训练设置成功", "model_type": request.model_type}
    except Exception as e:
        raise HTTPException(status_code=400, detail=str(e))

@router.post("/optimize/{model_type}")
async def run_model_optimization(model_type: str, rounds: int = 10):
    """运行模型优化"""
    try:
        model_enum = ModelType(model_type)
        result = await federated_integration.run_model_optimization(model_enum, rounds)
        return result
    except ValueError:
        raise HTTPException(status_code=400, detail=f"不支持的模型类型: {model_type}")
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@router.get("/model/{model_type}")
async def get_optimized_model(model_type: str):
    """获取优化后的模型"""
    try:
        model_enum = ModelType(model_type)
        model = federated_integration.get_optimized_model(model_enum)
        return {"model_type": model_type, "status": "available"}
    except ValueError:
        raise HTTPException(status_code=400, detail=f"不支持的模型类型: {model_type}")
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))